/////////////////////////
// File: camera_array_matching.cpp
// Authors: Brian Fehrman and Scott Logan and Matt Richard
// Date: 4-9-2013
// 
// Purpose: Source file for the matching functions
//
/////////////////////////

#include "camera_array_matching.h"

void template_matching_stereo()
{
   
}

//Use cam_1_1 as anchor
void template_matching_four_camera(  )
{
   string base_dir;
   string suffix_dir;
   FileStorage fs( "mat_save.xml", FileStorage::READ  );

   Mat gray_0_0, gray_0_1, gray_1_0, gray_1_1;
   Mat matches_0, matches_1, matches_2;
   Mat diff_0, diff_1, diff_2;
   Mat average_depth;
   Mat conf_0, conf_1, conf_2;
   
   Mat cam_0_0_img = imread( "images_new_array_16_sub/cam_0_0/31_gray_320_240.jpg" );
   Mat cam_0_1_img = imread( "images_new_array_16_sub/cam_0_1/31_gray_320_240.jpg" );
   Mat cam_1_0_img = imread( "images_new_array_16_sub/cam_1_0/31_gray_320_240.jpg"  );
   Mat cam_1_1_img = imread( "images_new_array_16_sub/cam_1_1/31_gray_320_240.jpg"  );
   
   cvtColor( cam_0_0_img, gray_0_0, CV_BGR2GRAY );
   cvtColor( cam_0_1_img, gray_0_1, CV_BGR2GRAY );
   cvtColor( cam_1_0_img, gray_1_0, CV_BGR2GRAY );
   cvtColor( cam_1_1_img, gray_1_1, CV_BGR2GRAY );
   
   gray_0_0.convertTo( gray_0_0, CV_32F, 1/255.0 );
   gray_0_1.convertTo( gray_0_1, CV_32F, 1/255.0 );
   gray_1_0.convertTo( gray_1_0, CV_32F, 1/255.0 );
   gray_1_1.convertTo( gray_1_1, CV_32F, 1/255.0 );
   
   /*
   Mat conf_0[] = { Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ),
                     Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ), 
                     Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ) 
                  };
                     
   Mat conf_1[] = { Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ),
                     Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ),
                     Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ), 
                  };
                  
   Mat conf_2[] = { Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ),
                     Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ),
                     Mat::zeros( gray_0_0.rows, gray_0_0.cols, CV_32F ), 
                  };
                  
   template_matching( gray_1_1, gray_1_0, 15, Point( 100, 15) , matches_0);
   cout << "Match 0 done" << endl;
   compute_diff_mat(  gray_1_1, matches_0, diff_0, true, false );
   
   fs << "diff_0" << diff_0;
   split( matches_0, conf_0 );
   fs << "conf_0" << conf_0[ 2 ];
   
   template_matching( gray_1_1, gray_0_1, 15, Point( 15, 100) , matches_1);
   cout << "Match 1 done" << endl;
   compute_diff_mat(  gray_1_1, matches_1, diff_1, false, true );
   
   fs << "diff_1" << diff_1;
   split( matches_1, conf_1 );
   fs << "conf_1" << conf_1[ 2 ];
   
   template_matching( gray_1_1, gray_0_0, 15, Point( 100, 100) , matches_2);
   cout << "Match 2 done" << endl;
   compute_diff_mat(  gray_1_1, matches_2, diff_2, true, true );
   
   fs << "diff_2" << diff_2;
   split( matches_2, conf_2 );
   fs << "conf_2" << conf_2[ 2 ];
   */
   
   fs[ "diff_0" ] >> diff_0;
   fs[ "diff_1" ] >> diff_1;
   fs[ "diff_2" ] >> diff_2;
   fs[ "conf_0" ] >> conf_0;
   fs[ "conf_1" ] >> conf_1;
   fs[ "conf_2" ] >> conf_2;

   average_depth = Mat::zeros( diff_0.rows, diff_0.cols, CV_8UC1 );
   
   for( int row = 0; row < diff_0.rows; row++ )
   {
      for( int col = 0; col < diff_0.cols; col++ )
      {
         
         float avg = 0.0;
         float std = 0.0;
         int diff_vals[ 3 ] = { 0 };
         float conf_vals[ 3 ] = { 0 };
         int num_used = 0;
         vector<float> vals_used;
         vector<float> conf_used;
         float max_conf = -1;
         int max_conf_idx = 0;
         float total_conf = 0.0;
         
         diff_vals[ 0 ] = diff_0.at< unsigned char >( row, col );
         diff_vals[ 1 ] = diff_1.at< unsigned char >( row, col );
         diff_vals[ 2 ] = diff_2.at< unsigned char >( row, col ) / sqrt( 2.0 ); //scale diagonal pair
         
         conf_vals[ 0 ] = conf_0.at< float >( row, col );
         conf_vals[ 1 ] = conf_1.at< float >( row, col );
         conf_vals[ 2 ] = conf_2.at< float >( row, col );
         
         for( int idx = 0; idx < 3; idx++ )
         {
            if( conf_vals[ idx ] > 0 )
            {
               vals_used.push_back( diff_vals[ idx ] );
               total_conf += conf_vals[ idx ];
               conf_used.push_back( conf_vals[ idx ] );
            
               if( conf_vals[ idx ] > max_conf )
               {
                  max_conf = conf_vals[ idx ];
                  max_conf_idx = vals_used.size() - 1;
               }
            }
         }
         
         for( int idx = 0; idx < vals_used.size(); idx++ )
         {
            avg += vals_used[ idx ];
         }
         
         avg /= vals_used.size() + 0.5;
         
         for( int idx = 0; idx < vals_used.size(); idx++ )
         {
            std += ( vals_used[ idx ] - avg ) * ( vals_used[ idx ] - avg );
         }
         
         std /= vals_used.size();
         
         std = sqrt( std );
         
         if( std > 12  )
         {
            //average_depth.at< unsigned char >( row, col ) = vals_used[ max_conf_idx ];
            average_depth.at< unsigned char >( row, col ) = *max_element( vals_used.begin(), vals_used.end() );
         }
         else
         {
            avg = 0.0;
            for( int idx = 0; idx < vals_used.size(); idx++ )
            {
               avg += vals_used[ idx ] * conf_used[ idx ] / total_conf ;
            }
            
            average_depth.at< unsigned char >( row, col ) = avg;
         }
         
      }
   }
   
   medianBlur( average_depth, average_depth, 7);
   
   namedWindow( "Image_1_1", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
   namedWindow( "Averaged", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
    namedWindow( "diff_0", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
     namedWindow( "diff_1", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
      namedWindow( "diff_2", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
   
   Mat diff_mat_color( average_depth.rows, average_depth.cols + 50, CV_8UC3 );
   Mat diff_mat_col_0( average_depth.rows, average_depth.cols + 50, CV_8UC3 );
   Mat diff_mat_col_1( average_depth.rows, average_depth.cols + 50, CV_8UC3 );
   Mat diff_mat_col_2( average_depth.rows, average_depth.cols + 50, CV_8UC3 );
   
   //cvtColor( diff_mat, diff_mat_color, CV_GRAY2RGB );
   //cvtColor( diff_mat_color, diff_mat_color, CV_RGB2HSV );
   
   
   for( int row = 0; row < average_depth.rows; row++ )
   {
      for( int col = 0; col < average_depth.cols; col++ )
      {
         diff_mat_color.at< Vec3b >( row, col )[ 0 ] = average_depth.at< unsigned char >( row, col ) * 1.5;
         diff_mat_color.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_color.at< Vec3b >( row, col )[ 2 ] = 255;
         
         diff_mat_col_0.at< Vec3b >( row, col )[ 0 ] = diff_0.at< unsigned char >( row, col ) * 1.5;
         diff_mat_col_0.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_col_0.at< Vec3b >( row, col )[ 2 ] = 255;
         
         diff_mat_col_1.at< Vec3b >( row, col )[ 0 ] = diff_1.at< unsigned char >( row, col ) * 1.5;
         diff_mat_col_1.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_col_1.at< Vec3b >( row, col )[ 2 ] = 255;
         
         diff_mat_col_2.at< Vec3b >( row, col )[ 0 ] = diff_2.at< unsigned char >( row, col ) * 1.5;
         diff_mat_col_2.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_col_2.at< Vec3b >( row, col )[ 2 ] = 255;
      }
   }
   
   for( int row =  0; row < average_depth.rows; row++ )
   {
      for( int col = average_depth.cols; col < average_depth.cols + 50; col++ )
      {
         diff_mat_color.at< Vec3b >( row, col )[ 0 ] =  ( float ) row / average_depth.rows * 150;
         diff_mat_color.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_color.at< Vec3b >( row, col )[ 2 ] = 255;
         
         diff_mat_col_0.at< Vec3b >( row, col )[ 0 ] =  ( float ) row / average_depth.rows * 150;
         diff_mat_col_0.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_col_0.at< Vec3b >( row, col )[ 2 ] = 255;
         
         diff_mat_col_1.at< Vec3b >( row, col )[ 0 ] =  ( float ) row / average_depth.rows * 150;
         diff_mat_col_1.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_col_1.at< Vec3b >( row, col )[ 2 ] = 255;
         
         diff_mat_col_2.at< Vec3b >( row, col )[ 0 ] =  ( float ) row / average_depth.rows * 150;
         diff_mat_col_2.at< Vec3b >( row, col )[ 1 ] = 255;
         diff_mat_col_2.at< Vec3b >( row, col )[ 2 ] = 255;
      }
   }
   
   cvtColor( diff_mat_color, diff_mat_color, CV_HSV2RGB );
   cvtColor( diff_mat_col_0, diff_mat_col_0, CV_HSV2RGB );
   cvtColor( diff_mat_col_1, diff_mat_col_1, CV_HSV2RGB );
   cvtColor( diff_mat_col_2, diff_mat_col_2, CV_HSV2RGB );

   imwrite( "Averaged.jpg", diff_mat_color );
   imwrite( "diff_0.jpg", diff_mat_col_0 );
   imwrite( "diff_1.jpg", diff_mat_col_1 );
   imwrite( "diff_2.jpg", diff_mat_col_2 );
   imwrite( "Image_1_1.jpg", gray_1_1 );
   
   for( ;; )
   {
      imshow( "Averaged", diff_mat_color );
      imshow( "diff_0", diff_mat_col_0 );
      imshow( "diff_1", diff_mat_col_1 );
      imshow( "diff_2", diff_mat_col_2 );
      imshow( "Image_1_1", gray_1_1 );
      
      if( waitKey( 30 ) > 0 ) break;
   }
   
}

//Assumes img is from left image and matches were found in right image
void compute_diff_mat( Mat img, Mat matches, Mat& diff_mat, bool x_change, bool y_change )
{
   int nrows = img.rows;
   int ncols = img.cols;
   double x_diff, y_diff;
   
   diff_mat = Mat::zeros( nrows, ncols, CV_8UC1 );
   
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {
         if( matches.at< Vec3f >( row, col )[ 2 ] > 0 )
         {
            x_diff = matches.at< Vec3f >( row, col )[ 0 ] - col; 
            y_diff = matches.at< Vec3f >( row, col )[ 1 ] - row;
   
            if( x_diff < 0 )
            {
               x_diff = 0;
            }
            if( y_diff < 0 )
            {
               y_diff = 0;
            }
            
            if( x_change && y_change )
            {
               diff_mat.at< unsigned char >( row, col ) = sqrt( x_diff * x_diff + y_diff * y_diff );
            }
            else if( y_change )
            {
               diff_mat.at< unsigned char >( row, col ) = y_diff;
            }
            else
            {
               diff_mat.at< unsigned char >( row, col ) = x_diff;
            }
         }
      }
   }
}

void template_matching( Mat left_img, Mat right_img, int template_size, Point search_region, Mat& matches )
{
   int t_size_div_2 = template_size / 2;
   int nrows = left_img.rows;
   int ncols = left_img.cols;
   matches = Mat::zeros( nrows, ncols, CV_32FC3 );
   double hit_thresh = 0.85;
   double max_thresh = 0.90;
   double proc_total = ( nrows - template_size );
   int proc_curr = 0;
   
   //GaussianBlur( left_img, left_img, Size( 3, 3 ), 1 , 0 );
   //GaussianBlur( img_2, img_2, Size( 3, 3 ), 1 , 0 );
   cout << fixed << setprecision( 6 );
   
   #pragma omp parallel for
   for( int row = 0; row < nrows - template_size; row++ )
   {
      for( int col = 0; col < ncols - template_size; col++ )
      {
         Mat des_temp;
         Mat search_temp;
         double max_response = -10000;
         double response = 0.0;
         double mean_1;
         double std_dev_1;
         double mean_2;
         double std_dev_2;
         Point max_loc;
         int num_hits = 0, num_tries = 0;
         
         //Create the desired template and compute its mean and stander deviation
         left_img.rowRange( row, row + template_size - 1 ).colRange( col, col + template_size - 1 ).copyTo( des_temp );
         mean_1 = compute_mean( des_temp );
         std_dev_1 = compute_std_dev( des_temp, mean_1 );
         
         for( int search_y = row; ( search_y < ( nrows - template_size ) ) && ( search_y <  ( row + search_region.y ) ); search_y++ )
         {
            for( int search_x = col; ( search_x < ( ncols - template_size ) ) && ( search_x < ( col + search_region.x ) ); search_x++ )
            {
               right_img.rowRange( search_y, search_y + template_size - 1 ).colRange( search_x, search_x + template_size - 1 ).copyTo( search_temp );
               mean_2 = compute_mean( search_temp );
               std_dev_2 = compute_std_dev( search_temp, mean_2 );
            
               response = normalized_correlation( des_temp, search_temp, mean_1, mean_2, std_dev_1, std_dev_2 );
               if( response > hit_thresh )
               {
                  num_hits++;
                  
                  if( response > max_response )
                  {
                     max_response = response;
                     max_loc.x = search_x + t_size_div_2;
                     max_loc.y = search_y + t_size_div_2;
                  }
               }
               
               num_tries++;
               
               if(response == 1) break;
            }
            
            if(response == 1) break;
         }
         if( max_response > max_thresh && ( ( (float) num_hits / num_tries ) < 0.5 ) )
         {
            matches.at< Vec3f >( row + t_size_div_2, col + t_size_div_2 )[ 0 ] = max_loc.x;
            matches.at< Vec3f >( row + t_size_div_2, col + t_size_div_2 )[ 1 ] = max_loc.y;
            matches.at< Vec3f >( row + t_size_div_2, col + t_size_div_2 )[ 2 ] = max_response;
         }
      }
      
       #pragma omp critical(dataupdate)
       {
         proc_curr++;
         cout << "Total processed: " << proc_curr / proc_total << "%" << endl;
       }
   }
}

//Compute the mean of a matrix
double compute_mean( Mat& img )
{
   int nrows = img.rows;
   int ncols = img.cols;
   double mean = 0.0;
   
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {
         mean += img.at< float >( row, col );
      }
   }
   
   mean /= ( nrows * ncols );
   
   return mean;
}

double compute_std_dev( Mat& img, double mean )
{
   int nrows = img.rows;
   int ncols = img.cols;
   double std_dev = 0.0;
   
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {
         double diff = img.at< float  >( row, col ) - mean;
         std_dev += ( diff * diff );
      }
   }
   
   std_dev /= ( nrows * ncols );
   
   std_dev = sqrt( std_dev );
   
   return std_dev;
}

double normalized_correlation( Mat& img_1, Mat& img_2, double mean_1, double mean_2,
                                 double std_dev_1, double std_dev_2 )
{
   int nrows = img_1.rows;
   int ncols = img_1.cols;
   double n_correlation = 0.0;   
   double num = 0, den = 0;
   double std_prod = ( std_dev_1 * std_dev_1 ) * ( std_dev_2 * std_dev_2 ) * ( nrows * ncols ) * ( nrows * ncols ); 
      
   for( int row = 0; row < nrows; row++ )
   {
      for( int col = 0; col < ncols; col++ )
      {
         double diff_1 = img_1.at< float  >( row, col ) - mean_1;
         double diff_2 = img_2.at< float >( row, col ) - mean_2;
         num += ( diff_1 * diff_2 );
      }
   }
   
   den = sqrt( std_prod );
   
   n_correlation = num / den;
   
   return n_correlation;
}
